import numpy as np
import open3d as o3d
from matplotlib import cm

# 文件路径设置（根据你需要保留哪一个）
data_dir = '202504\\3dponits\\office_11\\'
result_file = '202504\\3dponits\\Area_5-office_11_pred.npy'

# 加载数据
color = np.load(data_dir + 'color.npy')  # RGB颜色 (N,3) 0-255
coord = np.load(data_dir + 'coord.npy')  # 3D坐标 (N,3)
pred = np.load(result_file)              # 预测标签 (N,)

# 确保数据形状正确
assert len(coord) == len(color) == len(pred), "数据长度不匹配"

def visualize_with_open3d(coord, color, pred, downsample_factor=10):
    """
    使用Open3D可视化点云，其中原始点云和预测点云使用相同的颜色映射
    """
    # 下采样
    N = coord.shape[0]
    sample_idx = np.arange(0, N, downsample_factor)
    coord = coord[sample_idx]
    color = color[sample_idx]
    pred = pred[sample_idx]
    
    # 获取唯一标签并创建一致的颜色映射
    unique_labels = np.unique(pred)
    colormap = cm.get_cmap('tab20', len(unique_labels))
    label2color = {label: colormap(i)[:3] for i, label in enumerate(unique_labels)}  # 映射：label -> RGB颜色

    # 根据预测标签为所有点赋予统一的颜色
    mapped_colors = np.array([label2color[label] for label in pred])
    
    # 1. 原始点云（使用预测标签颜色显示）
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coord)
    pcd.colors = o3d.utility.Vector3dVector(mapped_colors)
    o3d.visualization.draw_geometries([pcd], window_name="原始点云(统一颜色映射)")

    # 2. 预测结果（也使用相同的颜色映射）
    pcd_pred = o3d.geometry.PointCloud()
    pcd_pred.points = o3d.utility.Vector3dVector(coord)
    pcd_pred.colors = o3d.utility.Vector3dVector(mapped_colors)
    o3d.visualization.draw_geometries([pcd_pred], window_name="预测分割结果(统一颜色映射)")

    # 3. 可选 - 边界点检测
    boundary_mask = np.zeros_like(pred, dtype=bool)
    for label in unique_labels:
        label_mask = (pred == label)
        # 简单边界检测（邻域变化）
        grad = np.abs(np.gradient(label_mask.astype(float)))
        if isinstance(grad, list):  # numpy.gradient 可能返回 list
            grad = sum(grad)
        boundary_mask |= grad > 0

    # 创建边界点云
    boundary_pcd = o3d.geometry.PointCloud()
    boundary_pcd.points = o3d.utility.Vector3dVector(coord[boundary_mask])
    boundary_pcd.colors = o3d.utility.Vector3dVector(np.tile([1, 0, 0], (np.sum(boundary_mask), 1)))  # 红色

    # 可视化叠加
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name="预测边界叠加")
    
    # 添加点云
    pcd_display = o3d.geometry.PointCloud()
    pcd_display.points = o3d.utility.Vector3dVector(coord)
    pcd_display.colors = o3d.utility.Vector3dVector(mapped_colors * 0.7)  # 降低亮度

    vis.add_geometry(pcd_display)
    vis.add_geometry(boundary_pcd)

    # 设置渲染参数
    opt = vis.get_render_option()
    opt.point_size = 2.0
    opt.background_color = np.asarray([0.1, 0.1, 0.1])  # 深色背景

    vis.run()
    vis.destroy_window()

# 执行可视化
visualize_with_open3d(coord, color, pred, downsample_factor=5)
